#!/usr/bin/env python3
"""
Script to count words in DOCX files from multiple directories.
Only processes files starting with 'AA' and excludes subdirectories.
"""

import os
import sys
import argparse
import zipfile
import tempfile
import shutil
import xml.etree.ElementTree as ET
from docx import Document
from docx.shared import Pt
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.enum.dml import MSO_THEME_COLOR_INDEX
import warnings


def extract_text_from_xml(filepath):
    """
    Extract text directly from DOCX XML when python-docx fails.
    
    Args:
        filepath (str): Path to the DOCX file
        
    Returns:
        list: List of text paragraphs, or None if extraction failed
    """
    try:
        temp_dir = tempfile.mkdtemp()
        paragraphs = []
        
        # Extract the DOCX file
        with zipfile.ZipFile(filepath, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        # Read the document.xml file
        doc_xml_path = os.path.join(temp_dir, 'word', 'document.xml')
        if not os.path.exists(doc_xml_path):
            shutil.rmtree(temp_dir)
            return None
        
        # Parse the XML
        tree = ET.parse(doc_xml_path)
        root = tree.getroot()
        
        # Define namespaces
        namespaces = {
            'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'
        }
        
        # Find all paragraph elements (w:p)
        paragraph_elements = root.findall('.//w:p', namespaces)
        
        for para_elem in paragraph_elements:
            # Find all text elements within this paragraph
            text_elements = para_elem.findall('.//w:t', namespaces)
            
            # Combine all text in this paragraph
            paragraph_text = ""
            for text_elem in text_elements:
                if text_elem.text:
                    paragraph_text += text_elem.text
            
            # Add paragraph if it has content
            if paragraph_text.strip():
                paragraphs.append(paragraph_text.strip())
        
        # Clean up
        shutil.rmtree(temp_dir)
        
        return paragraphs if paragraphs else None
        
    except Exception as e:
        print(f"    XML extraction failed: {str(e)}")
        if 'temp_dir' in locals():
            shutil.rmtree(temp_dir, ignore_errors=True)
        return None


def repair_docx_file(filepath):
    """
    Attempt to repair a corrupted DOCX file by extracting and re-packaging.
    
    Args:
        filepath (str): Path to the corrupted DOCX file
        
    Returns:
        str: Path to repaired file, or None if repair failed
    """
    try:
        # Create a temporary directory
        temp_dir = tempfile.mkdtemp()
        
        # Extract the DOCX file (it's a zip archive)
        with zipfile.ZipFile(filepath, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        # Check if the extracted files contain the required DOCX structure
        required_files = ['word/document.xml', 'word/styles.xml', '[Content_Types].xml']
        missing_files = []
        
        for req_file in required_files:
            if not os.path.exists(os.path.join(temp_dir, req_file)):
                missing_files.append(req_file)
        
        if missing_files:
            print(f"    Missing required files: {', '.join(missing_files)}")
            shutil.rmtree(temp_dir)
            return None
        
        # Create a new repaired DOCX file
        repaired_path = filepath + '.repaired.docx'
        
        # Create a new zip file with the extracted contents
        with zipfile.ZipFile(repaired_path, 'w', zipfile.ZIP_DEFLATED) as zip_out:
            for root, dirs, files in os.walk(temp_dir):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, temp_dir)
                    zip_out.write(file_path, arcname)
        
        # Clean up
        shutil.rmtree(temp_dir)
        
        return repaired_path
        
    except Exception as e:
        print(f"    Repair attempt failed: {str(e)}")
        if 'temp_dir' in locals():
            shutil.rmtree(temp_dir, ignore_errors=True)
        return None


def get_aa_docx_files(directories):
    """
    Get all DOCX files starting with 'AA' from specified directories.
    Excludes subdirectories and temporary files.
    
    Args:
        directories (list): List of directory paths to search
        
    Returns:
        list: List of tuples (filepath, filename, directory)
    """
    all_files = []
    
    for directory in directories:
        if not os.path.isdir(directory):
            print(f"Warning: Directory '{directory}' does not exist, skipping...")
            continue
        
        print(f"Scanning directory: {os.path.abspath(directory)}")
        
        # Get all files in the directory (not subdirectories)
        try:
            files = os.listdir(directory)
        except PermissionError:
            print(f"Warning: Cannot access directory '{directory}', skipping...")
            continue
        
        aa_files = []
        for f in files:
            filepath = os.path.join(directory, f)
            
            # Only process files (not directories)
            if os.path.isfile(filepath):
                # Check if file starts with 'AA' and ends with '.docx'
                if f.startswith('AA') and f.lower().endswith('.docx'):
                    # Skip temporary files and repair files
                    if (not f.startswith('.~') and 
                        '~' not in f and 
                        not f.endswith('.repaired.docx')):
                        aa_files.append((filepath, f, directory))
        
        aa_files.sort()  # Sort files alphabetically
        all_files.extend(aa_files)
        
        if aa_files:
            print(f"  Found {len(aa_files)} files starting with 'AA':")
            for i, (_, filename, _) in enumerate(aa_files, 1):
                print(f"    {i}. {filename}")
        else:
            print(f"  No files starting with 'AA' found in this directory")
    
    return all_files


def count_words_in_text(text):
    """
    Count words in text, handling Chinese and English characters properly.
    For Chinese: count each Chinese character as one word
    For English: count space-separated words
    For mixed content: use appropriate counting for each part
    
    Args:
        text (str): Text to count words in
        
    Returns:
        int: Word count
    """
    if not text or not text.strip():
        return 0
    
    text = text.strip()
    word_count = 0
    
    # Process the text character by character
    i = 0
    n = len(text)
    
    while i < n:
        char = text[i]
        
        # Check if it's a Chinese character (including common punctuation)
        if '\u4e00' <= char <= '\u9fff':  # CJK Unified Ideographs
            word_count += 1
            i += 1
        # Check if it's Chinese punctuation
        elif char in '，。、；：？！「」『』（）【】《》':
            word_count += 1
            i += 1
        # Check if it's whitespace
        elif char.isspace():
            i += 1
        # Otherwise, treat as English/other text and find word boundaries
        else:
            # Find the next whitespace or Chinese character
            word_start = i
            while i < n and not text[i].isspace() and not ('\u4e00' <= text[i] <= '\u9fff'):
                i += 1
            
            # If we found a non-empty word, count it
            if i > word_start:
                word_count += 1
    
    return word_count


def count_docx_words(files_list):
    """
    Count words in DOCX files from files list.
    Handles corrupted files with repair and XML extraction.
    
    Args:
        files_list (list): List of tuples (filepath, filename, directory)
    """
    if not files_list:
        print("No files to count!")
        return
    
    print(f"\nFound {len(files_list)} files to count:")
    for i, (_, filename, directory) in enumerate(files_list, 1):
        print(f"  {i}. {filename} (from: {directory})")
    
    total_words = 0
    successful_counts = 0
    file_word_counts = []
    total_pages = 0
    file_page_counts = []
    
    # Count words in each document
    for i, (filepath, filename, directory) in enumerate(files_list, 1):
        print(f"\nCounting words in {filename} ({i}/{len(files_list)})...")
        print(f"  From: {directory}")
        
        # Try to open the file normally first
        src_doc = None
        repaired_file = None
        original_filepath = filepath
        word_count = 0
        page_count = 0
        
        try:
            # Test if file can be opened
            test_doc = Document(filepath)
            print(f"  ✓ File is readable, proceeding with word count...")
            
            # Open the source document again for actual processing
            src_doc = Document(filepath)
        except Exception as e:
            error_msg = str(e)
            if "NULL" in error_msg or "archive" in error_msg:
                print(f"  ✗ Error processing {filename}: Corrupted DOCX file structure")
                print(f"    Attempting to repair...")
                
                # Try to repair the file
                repaired_file = repair_docx_file(filepath)
                if repaired_file:
                    print(f"  ✓ File repaired successfully, using repaired version")
                    try:
                        src_doc = Document(repaired_file)
                        filepath = repaired_file  # Use repaired file for processing
                    except Exception as repair_error:
                        print(f"  ✗ Repaired file still cannot be opened: {str(repair_error)}")
                        src_doc = None
                
                # If repair failed or repaired file still can't be opened, try XML extraction
                if src_doc is None:
                    print(f"  ✗ Repair failed, trying XML text extraction...")
                    
                    # Try to extract text directly from XML
                    text_paragraphs = extract_text_from_xml(original_filepath)
                    if text_paragraphs:
                        print(f"  ✓ XML extraction successful, counting words from extracted text")
                        # Count words from extracted text
                        for para_text in text_paragraphs:
                            word_count += count_words_in_text(para_text)
                        src_doc = None  # Skip normal processing
                    else:
                        print(f"  ✗ XML extraction failed, skipping this file")
                        src_doc = None
            else:
                print(f"  ✗ Error processing {filename}: {error_msg}")
                src_doc = None
        
        if src_doc is None and word_count == 0:
            print(f"    Skipping this file and continuing...")
            # Clean up repaired file if it exists
            if repaired_file and os.path.exists(repaired_file):
                os.remove(repaired_file)
            continue
        
        # If we have a normal document, count words from it
        if src_doc is not None:
            try:
                # Suppress warnings during processing
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    
                    # Count words from paragraphs
                    para_count = 0
                    for paragraph in src_doc.paragraphs:
                        para_count += 1
                        if para_count % 1000 == 0:
                            print(f"    Processing paragraph {para_count}/{len(src_doc.paragraphs)}...")
                        
                        word_count += count_words_in_text(paragraph.text)
                    
                    # Count words from tables
                    for table in src_doc.tables:
                        for row in table.rows:
                            for cell in row.cells:
                                for para in cell.paragraphs:
                                    word_count += count_words_in_text(para.text)
                
                print(f"  ✓ Successfully counted words in {filename} ({len(src_doc.paragraphs)} paragraphs, {len(src_doc.tables)} tables)")
                
            except Exception as e:
                error_msg = str(e)
                if "WD_COLOR_INDEX" in error_msg or "XML mapping" in error_msg:
                    print(f"  ✗ Error processing {filename}: Color formatting issue")
                    print(f"    This file has unsupported color formatting")
                elif "NULL" in error_msg or "archive" in error_msg:
                    print(f"  ✗ Error processing {filename}: Corrupted DOCX file structure")
                    print(f"    This file may be damaged or have formatting issues")
                else:
                    print(f"  ✗ Error processing {filename}: {error_msg}")
                print(f"    Skipping this file and continuing...")
                # Clean up repaired file if it exists
                if repaired_file and os.path.exists(repaired_file):
                    os.remove(repaired_file)
                continue
            finally:
                # Clean up repaired file if it exists
                if repaired_file and os.path.exists(repaired_file):
                    os.remove(repaired_file)
        
        total_words += word_count
        successful_counts += 1
        file_word_counts.append((filename, word_count))
        
        # Estimate page count with improved algorithm
        # For Chinese-heavy documents, use different ratios
        # Chinese characters take more space than English words
        if word_count <= 500:
            page_count = 1
        elif word_count <= 1500:
            page_count = 2
        elif word_count <= 3000:
            page_count = round(word_count / 800)
        elif word_count <= 10000:
            page_count = round(word_count / 600)
        else:
            # For larger documents, use more conservative estimate
            page_count = round(word_count / 500)
        total_pages += page_count
        file_page_counts.append((filename, page_count))
        
        print(f"  Word count: {word_count:,} words")
        print(f"  Page count: {page_count} pages")
    
    # Print individual file counts
    if file_word_counts:
        print(f"\n" + "=" * 50)
        print(f"INDIVIDUAL FILE COUNTS")
        print(f"=" * 50)
        print(f"{'Filename':<30} {'Words':>12} {'Pages':>8}")
        print("-" * 52)
        for filename, word_count in sorted(file_word_counts, key=lambda x: x[1], reverse=True):
            page_count = next((p for f, p in file_page_counts if f == filename), 0)
            print(f"{filename:<30} {word_count:>12,} {page_count:>8}")
    
    # Print summary at the end
    print(f"\n" + "=" * 50)
    print(f"SUMMARY")
    print(f"=" * 50)
    print(f"Files successfully processed: {successful_counts}/{len(files_list)}")
    if successful_counts < len(files_list):
        print(f"Files skipped due to errors: {len(files_list) - successful_counts}")
    print(f"Total word count: {total_words:,} words")
    print(f"Total page count: {total_pages} pages")


def main():
    parser = argparse.ArgumentParser(
        description='Count words in DOCX files starting with "AA" from multiple directories',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
Examples:
  python3 count_docx.py dir1 dir2 dir3
  python3 count_docx.py /path/to/docs /path/to/more/docs
        '''
    )
    parser.add_argument('directories', nargs='+', 
                       help='Directories to search for AA*.docx files')
    
    args = parser.parse_args()
    
    # Validate directories
    valid_dirs = []
    for directory in args.directories:
        if os.path.isdir(directory):
            valid_dirs.append(directory)
        else:
            print(f"Warning: Directory '{directory}' does not exist and will be skipped")
    
    if not valid_dirs:
        print("Error: No valid directories specified")
        sys.exit(1)
    
    print("DOCX Word Counter - AA Files Only")
    print("=" * 50)
    print(f"Searching for AA*.docx files in:")
    for i, directory in enumerate(valid_dirs, 1):
        print(f"  {i}. {os.path.abspath(directory)}")
    print("-" * 50)
    
    # Get all AA docx files
    files_list = get_aa_docx_files(valid_dirs)
    
    if not files_list:
        print("No files starting with 'AA' found in any of the specified directories")
        sys.exit(1)
    
    # Count words in the files
    count_docx_words(files_list)


if __name__ == "__main__":
    main()